/* Copyright (C) 2005, 2006, 2007, 2008, 2009, 2010 Mario Orsi
   This file is part of Brahms.
   Brahms is free software: you can redistribute it and/or modify it under the terms of the GNU General Public License as published by
   the Free Software Foundation, either version 3 of the License, or (at your option) any later version.
   Brahms is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License for more details.
   You should have received a copy of the GNU General Public License along with Brahms.  If not, see <http://www.gnu.org/licenses/>. */

/*************************************************************************************************************************************
 * stepInt.c -- Functions for integratig the equations of motion at every step, and for optional control of pressure and temperature *
 *************************************************************************************************************************************/

#include "dataStructs.h"

extern VecR *siteCoords, *siteVelocity, *siteOrientVec, *siteAngMom, *siteForce, *siteTorque; 
extern RMat *siteRotMat;
extern const int DOFs, flexBox, keepSquare, keepTetragonal, nSites, nSolutes, nWaters, stepCount, zConstraint;
extern const real tauP, tauT, deltaT, extPressure, extTemperature, *mass, *inertia, totalMass;
extern const LipidStruct dopcStruct, dopeStruct;
extern Prop boxVolume, pressure, xyArea;
extern RMat boxMatrix;
extern Site *site;  
extern const VecR solInertia;
extern VecR region;

// external function declarations:
void BuildRxMatrix( RMat*, real ); 
void BuildRyMatrix( RMat*, real ); 
void BuildRzMatrix( RMat*, real ); 
real CalcKinEnergy();
void CalcPressTens( RMat*, RMat );
real CalcPressure( real );
void GetOrientPrincAxisZ_FromRotMat( VecR*, RMat );
void UpdateQMatrix( RMat*, RMat* );

void ApplyThermostat() // weak-coupling scheme [Berendsen et al, J Chem Phys, 81, 3684 (1984)]
{ 
  const int nRotDofSymmRigBody = 2; // number of rotational degrees of freedom of a symmetric rigid body
  const int nRotDofNonSymmRigBody = 3; // number of rotational degrees of freedom of a non-symmetric rigid body
  const int nTransDofAnyBody = 3; // number of translational degrees of freedom of any body
  int n; // counter used in DO_SITE
  int nRotDofSymmRigBodyLipids = 0; // total number of rotational degrees of freedom from all symmetric rigid bodies in the system
  int nTransDofAnyBodyLipids = 0; // total number of translational degrees of freedom in the system
  real lipTransTemp, lipRotTemp, watTransTemp, watRotTemp, soluteTransTemp, soluteRotTemp; // translational and rotational temperatures
  real ltScalingFac, lrScalingFac, wtScalingFac, wrScalingFac, stScalingFac, srScalingFac; 
  real rotSolKinEn = 0., transSolKinEn = 0., rotWatKinEn = 0., transWatKinEn = 0., rotLipKinEn = 0., transLipKinEn = 0.; 
  
  DO_SITE { // accumulating kinetic energy components
    if ( WATER_TYPE == site[ n ].type ) {
      rotWatKinEn   += 0.5 * VLenSq( siteAngMom[ n ] ) / inertia[ site[ n ].type ];              
      transWatKinEn += 0.5 * VLenSq( siteVelocity[ n ] ) * mass[ site[ n ].type ]; 
    } else if  ( SOLUTE == site[ n ].type ) {
      rotSolKinEn   += 0.5 * VSumSqDiv( siteAngMom[ n ], solInertia );              
      transSolKinEn += 0.5 * VLenSq( siteVelocity[ n ] ) * mass[ site[ n ].type ];
    } else { // if not water
      transLipKinEn += 0.5 * VLenSq( siteVelocity[ n ] ) * mass[ site[ n ].type ]; 
      nTransDofAnyBodyLipids += nTransDofAnyBody;
      if ( SRB == site[ n ].mechType ) { // rotational contribution from symmetric rigid body sites
	rotLipKinEn += 0.5 * VLenSq( siteAngMom[ n ] ) / inertia[ site[ n ].type ]; 
	nRotDofSymmRigBodyLipids += nRotDofSymmRigBody;
      }
    }
  }
  
  if ( zConstraint ) {
    soluteRotTemp   = 2. * rotSolKinEn   / ( nRotDofNonSymmRigBody ); // assuming only 1 solute
    soluteTransTemp = 2. * transSolKinEn / ( nTransDofAnyBody - 1 ); // 1 DOF is z-constrained
  } else {
    soluteRotTemp   = 2. * rotSolKinEn   / ( nRotDofNonSymmRigBody * nSolutes ); 
    soluteTransTemp = 2. * transSolKinEn / ( nTransDofAnyBody * nSolutes ); 
  }
  
  watRotTemp    = 2. * rotWatKinEn   / ( kB_IN_BRAHMS_UNITS * nRotDofSymmRigBody * nWaters ); 
  watTransTemp  = 2. * transWatKinEn / ( kB_IN_BRAHMS_UNITS * nTransDofAnyBody * nWaters ); 

  lipRotTemp = 2. * rotLipKinEn   / ( kB_IN_BRAHMS_UNITS * nRotDofSymmRigBodyLipids ); 
  lipTransTemp = 2. * transLipKinEn / ( kB_IN_BRAHMS_UNITS * nTransDofAnyBodyLipids ); 
    
  ltScalingFac = sqrt( 1 + deltaT / tauT * ( ( extTemperature / lipTransTemp ) - 1. ) ); 
  lrScalingFac = sqrt( 1 + deltaT / tauT * ( ( extTemperature / lipRotTemp )   - 1. ) ); 
  wtScalingFac = sqrt( 1 + deltaT / tauT * ( ( extTemperature / watTransTemp ) - 1. ) ); 
  wrScalingFac = sqrt( 1 + deltaT / tauT * ( ( extTemperature / watRotTemp ) - 1. ) ); 
  stScalingFac = sqrt( 1 + deltaT / ( 0.2 * tauT ) * ( ( extTemperature / soluteTransTemp ) - 1. ) ); 
  srScalingFac = sqrt( 1 + deltaT / ( 0.2 * tauT ) * ( ( extTemperature / soluteRotTemp ) - 1. ) ); 
   
  DO_SITE { 
    if ( WATER_TYPE == site[ n ].type ) {
      VScale( siteVelocity[ n ], wtScalingFac ); 
      VScale( siteAngMom[ n ], wrScalingFac );
    } else if ( SOLUTE == site[ n ].type ) {
      VScale( siteVelocity[ n ], stScalingFac ); 
      VScale( siteAngMom[ n ], srScalingFac );
    } else {
      VScale( siteVelocity[ n ], ltScalingFac );
      if ( SRB == site[ n ].mechType ) { 
	VScale( siteAngMom[ n ], lrScalingFac );
      }
    }
  }
}

void ApplyBarostat( const RMat pressureTensor ) // weak-coupling scheme [Berendsen et al, J Chem Phys, 81, 3684 (1984)]
{
  int n;
  real p;
  RMat deltaPress, identity, muTens, pTens;
  const real BETA = 4.6e-5 / atm_IN_BRAHMS_UNITS; // isothermal compressibility converted into Brahms units 

  identity.u[0] = identity.u[4] = identity.u[8] = 1.; // setting identity matrix as [100;010;001]
  identity.u[1] = identity.u[2] = identity.u[3] = identity.u[5] = identity.u[6] = identity.u[7] = 0.;
  MZero( pTens.u );

  // commenting out the following line results in considering P(t) instead of P(t+dt) 
  //  CalcPressTens( pressureTensor, virialTensor );  // evaluating pressure tensor
  
  // preparing pTens, the actual pressure tensor used in the algorithm
  if ( !flexBox ) { // isotropic NPT - pTens is diagonal and all elements are equal to trace(pressureTensor)/3
    p = Trace( pressureTensor.u ) / 3.;
    pTens.u[ 0 ] = pTens.u[ 4 ] = pTens.u[ 8 ] = p;
  } else {          // anisotropic NPT 
    if ( keepSquare ) { // semi-anisotropic NPT: constrain on xy to maintain a square
      p = ( pressureTensor.u[ 0 ] + pressureTensor.u[ 4 ] ) / 2.;
      pTens.u[ 0 ] = pTens.u[ 4 ] = p;
      pTens.u[ 8 ] = pressureTensor.u[ 8 ];
    } else {         // general case
      MCopy( pTens.u, pressureTensor.u );
    }
  }
  
  MSAdd( deltaPress.u, pTens.u, -extPressure, identity.u ); // pressTens - Pext * 1
  MSAdd( muTens.u, identity.u, deltaT * BETA / ( 3. * tauP ), deltaPress.u ); // compute mu tensor
  if ( keepTetragonal ) MZeroOffDiags( muTens.u ); // constrain non-diagonal elements to zero - NPxPyPzT ensemble
  boxMatrix.u[0] *= muTens.u[0];
  boxMatrix.u[4] *= muTens.u[4];
  boxMatrix.u[8] *= muTens.u[8];
  DO_SITE { // scale particles' coordinates
    siteCoords[ n ].x *= muTens.u[0]; 
    siteCoords[ n ].y *= muTens.u[4]; 
    siteCoords[ n ].z *= muTens.u[8]; 
  }
  
  region.x = boxMatrix.u[ 0 ]; // only for orthorombic geometry!!!!
  region.y = boxMatrix.u[ 4 ];
  region.z = boxMatrix.u[ 8 ];
  boxVolume.val = region.x * region.y * region.z;
  
  // boxVolume.val = MatDet( boxMatrix.u );
  //  boxVolume.val *= MatDet( muTens.u );
  xyArea.val = region.x * region.y;

  //  if (stepCount%stepAvg==0)  PrintMatrix(muTens);

  // isotropic NPT as a special case   
  //     currentPressure = CalcPressure( temp.val );
  //     scalingFac = 1. - deltaT * BETA / ( 3. * tauP ) * ( extPressure - currentPressure );
  //     //scalingFac = pow( scalingFac, 1./3. );
  //     // VScale( region, scalingFac );
  //     MScale( boxMatrix.u, scalingFac );
  //     DO_SITE VScale( siteCoords[ n ], scalingFac );
}

/******************************************************************************************
 * RotateRigidBody -- underlying algorithm from Dullweber et al. [J Chem Phys, 107, 5840] *
 ******************************************************************************************/
static void RotateRigidBody( const int n, RMat Q ) 
{
  real phi1, phi2, phi3, halfDeltaT = deltaT / 2.;
  VecR uOld, uNew, princInertia; // principal moments of inertias Ix Iy Iz     
  RMat R;

  if ( site[ n ].type == SOLUTE ) VCopy( princInertia, solInertia );
  else VSet( princInertia, inertia[ site[ n ].type ], inertia[ site[ n ].type ], 0 ); // symmetric rigid body

  VCopy( uNew, siteAngMom[ n ] );
  phi1 = halfDeltaT * uNew.x / princInertia.x;  // FIRST ROTATION
  BuildRxMatrix( &R, phi1 );                    // compute R1
  VCopy( uOld, uNew );                          // uOld:= uNew
  MVMul( uNew, R.u, uOld );                     // rotate angular momentum: uNew = R1 * uOld
  UpdateQMatrix( &Q, &R );                      // Q := Q * R^T
  if ( site[ n ].mechType == NSRB ) phi2 = halfDeltaT * uNew.y / princInertia.y; // SECOND ROTATION
  else phi2 = deltaT * uNew.y / princInertia.y; // SECOND ROTATION for symmetric body
  BuildRyMatrix( &R, phi2 );                    // compute R2
  VCopy( uOld, uNew );                          // uOld:= uNew
  MVMul( uNew, R.u, uOld );                     // rotate angular momentum: uNew = R2 * uOld
  UpdateQMatrix( &Q, &R );                      // Q := Q * R^T
  if ( site[ n ].mechType == NSRB ) { // general case: non-linear rigid body
    phi3 = deltaT * uNew.z / princInertia.z;      // THIRD ROTATION
    BuildRzMatrix( &R, phi3 );                    // compute R3
    VCopy( uOld, uNew );                          // uOld:= uNew
    MVMul( uNew, R.u, uOld );                     // rotate angular momentum: uNew = R3 * uOld
    UpdateQMatrix( &Q, &R );                      // Q := Q * R^T
    phi2 = halfDeltaT * uNew.y / princInertia.y;  // FOURTH ROTATION
    BuildRyMatrix( &R, phi2 );                    // compute R4
    VCopy( uOld, uNew );                          // uOld:= uNew
    MVMul( uNew, R.u, uOld );                     // rotate angular momentum: uNew = R4 * uOld
    UpdateQMatrix( &Q, &R );                      // Q := Q * R^T
  } // END if (  site[ n ].mechType == NSRB )
  phi1 = halfDeltaT * uNew.x / princInertia.x;  // FIFTH (OR THIRD) ROTATION
  BuildRxMatrix( &R, phi1 );                    // compute R5
  VCopy( uOld, uNew );                          // uOld:= uNew
  MVMul( uNew, R.u, uOld );                     // rotate angular momentum: uNew = R5 * uOld
  UpdateQMatrix( &Q, &R );                      // Q := Q * R^T
  siteRotMat[ n ] = Q;                  // now we have Q( t + dt )
  VCopy( siteAngMom[ n ], uNew );      // and h(t+dt )?
  GetOrientPrincAxisZ_FromRotMat( &siteOrientVec[ n ], Q ); // obtain orientation vector from rotation matrix
}

/************************************************
 * LeapfrogStep -- adapted from Rapaport [p.27] *
 ************************************************/
void LeapfrogStep( int part ) 
{
  int n;
  real halfDeltaT = deltaT / 2.;
  VecR bodyTorque;     
  RMat Q;
  if ( part == 1 ) { // Integration part-1  
    DO_SITE {
      VVSAdd( siteVelocity[n], halfDeltaT / mass[site[n].type], siteForce[n] ); // v(t) -> v(t+dt/2)
      VVSAdd( siteCoords[n], deltaT, siteVelocity[n] ); // advance mass centre coordinates from t to t+dt
      if ( site[ n ].mechType == NSRB || site[ n ].mechType == SRB ) { // if site 'n' has rotational degrees of freedom
	Q = siteRotMat[ n ];
	MVMul( bodyTorque, Q.u, siteTorque[ n ] ); // T^b = R * T^S
	VVSAdd( siteAngMom[ n ], halfDeltaT, bodyTorque ); // advance angular momenta from t to t+dt/2
	RotateRigidBody( n, Q ); // Q(t) -> Q(t+dt)
      } 
    } // end DO_SITE
  } else { // Integration part-2
    DO_SITE {
      VVSAdd( siteVelocity[ n ], halfDeltaT / mass[ site[ n ].type ], siteForce[ n ] ); // v(t+dt/2) -> v(t+dt)
      if ( site[ n ].mechType == NSRB || site[ n ].mechType == SRB ) { // if site 'n' has rotational degrees of freedom
	MVMul( bodyTorque, siteRotMat[ n ].u, siteTorque[ n ] ); // T^b = R * T^S
	VVSAdd( siteAngMom[ n ], halfDeltaT, bodyTorque ); // advance momenta in the body-fixed principal frame
      } // end if
    } // end DO_SITE
  } // end else part-2
} 
